// condaUtils
const fs = require('fs');
const path = require('path');
const axios = require('axios');
const { execSync } = require('child_process');

class CondaManager {
    constructor(app, remote_id, systemInfo) {
        this.app = app;
        this.remote_id = remote_id;
        this.systemInfo = systemInfo;
        this.installDirs = {
            conda: '$HOME/.labelapp/miniconda',
            node: '$HOME/.labelapp/nodejs'
        };
        this.condaPath = remote_id === 'local' ? 'conda' : `${this.installDirs.conda}/bin/conda`;
        this.pytorchVersions = this.loadPytorchVersions();
        this.paddleVersions = this.loadPaddleVersions();
        this.executeCommand = this.setupExecutor();
    }

    loadPytorchVersions() {
        const filePath = path.join(__dirname, 'configs', 'pytorch_conda_versions.json');
        const rawData = fs.readFileSync(filePath);
        return JSON.parse(rawData);
    }

    loadPaddleVersions() {
        const filePath = path.join(__dirname, 'configs', 'paddle_versions.json');
        const rawData = fs.readFileSync(filePath);
        return JSON.parse(rawData);
    }

    setupExecutor() {
        if (this.remote_id === 'local') {
            return (command) => {
                // console.log('execSync: ', command);
                try {
                    return execSync(command).toString();
                } catch (error) {
                    console.error(`Local command execution failed: ${command}\n`, error);
                    throw error;
                }
            };
        } else {
            return async (command) => {
                let commandEx = command;
                if (command.startsWith('conda')) {
                    commandEx = `${this.installDirs.conda}/bin/${command}`;
                }

                const result = await this.app.remotes.execute(this.remote_id, commandEx);
                if (result.code == 0) {
                    return result.msg;
                } else {
                    throw new Error(JSON.stringify(result));
                }
            };
        }
    }

    async getEnvList() {
        const command = `${this.condaPath} env list`;
        try {
            const envListResult = await this.executeCommand(command);
            const envList = this.parseEnvList(envListResult);
            return envList;
        } catch (error) {
            console.error('Error getting Conda environments:', error.message);
            return [];
        }
    }

    async getPipPackages(env) {
        // 确保 env 是字符串
        if (typeof env !== 'string') {
            throw new Error(`Invalid environment: ${JSON.stringify(env)}`);
        }

        const command = `${this.condaPath} run -n ${env} pip list --format=json 2>/dev/null`;
        try {
            const pipListResult = await this.executeCommand(command);
            return JSON.parse(pipListResult);
        } catch (error) {
            console.error(`Error getting pip packages for environment ${env}:`, error.message);
            return [];
        }
    }

    parseEnvList(output) {
        const lines = output.trim().split('\n');
        const environments = [];

        lines.forEach((line) => {
            if (line && !line.startsWith('#')) {
                const parts = line.trim().split(/\s+/);
                if (parts.length >= 2) {
                    let envName, envPath;

                    if (parts[1] === '*') {
                        envName = parts[0];
                        envPath = parts[2];
                    } else {
                        envName = parts[0];
                        envPath = parts[1];
                    }

                    environments.push({
                        name: envName,
                        path: envPath
                    });
                }
            }
        });

        return environments;
    }

    async createEnv(envName, envPath) {
        try {
            const scriptUrl = `https://label-anything.bd.bcebos.com/labelapp/conda/${envPath}`;
            const response = await axios.get(scriptUrl);
            const envScript = response.data.trim();

            const lines = envScript.split('\n');
            const pythonLine = lines.find((line) => line.startsWith('python'));
            const frameworkLine = lines.find((line) => line.startsWith('framework'));
            let command;

            // 解析Python版本
            if (pythonLine) {
                const pythonVersion = pythonLine.split(' ')[1];
                command = `${this.condaPath} create -y -n ${envName} python=${pythonVersion}`;
                await this.executeCommand(command);
                console.log(`Environment ${envName} created with Python ${pythonVersion}`);
            } else {
                throw Error('No python specified!');
            }

            // 解析框架版本
            if (frameworkLine) {
                const framework = frameworkLine.split(' ')[1];
                if (framework.startsWith('pytorch')) {
                    await this.installPyTorch(envName);
                } else if (framework.startsWith('paddle')) {
                    await this.installPaddle(envName, framework);
                } else {
                    throw Error(`Unknown framework ${framework}`);
                }
            } else {
                throw Error('No framework specified!');
            }

            // 安装其他包
            for (const line of lines) {
                if (line.startsWith('conda')) {
                    const cmd = line.replace('$ENV_NAME', envName);
                    await this.executeCommand(cmd);
                    console.log(`Executed: ${cmd}`);
                }
            }
            return true;
        } catch (error) {
            console.error('Error creating environment:', error);
            return null;
        }
    }

    getClosestVersionCommand(cudaVersion, envName) {
        const availableVersions = Object.keys(this.pytorchVersions).filter((key) => key.startsWith('CUDA'));
        const targetMajor = parseInt(cudaVersion.split('.')[0], 10);
        const targetMinor = parseFloat(cudaVersion);

        const matchingVersions = availableVersions.filter((v) => parseInt(v.split(' ')[1].split('.')[0], 10) === targetMajor);

        if (!matchingVersions.length) {
            return this.pytorchVersions['CPU Only'][0];
        }
        const closest = matchingVersions.reduce((prev, curr) => {
            const currMinor = parseFloat(curr.split(' ')[1]);
            return Math.abs(currMinor - targetMinor) < Math.abs(parseFloat(prev.split(' ')[1]) - targetMinor) ? curr : prev;
        });
        const command = this.pytorchVersions[closest][0] || this.pytorchVersions['CPU Only'][0];
        return `conda install -y -n ${envName} ${command.split(' ').slice(2).join(' ')}`;
    }

    async installPyTorch(envName) {
        try {
            let cudaVersion = this.systemInfo.cudaInfo.nvccVersion;
            if (!cudaVersion || cudaVersion === 'null') {
                cudaVersion = this.systemInfo.cudaInfo.nvidiaSmiCUDA;
            }

            let command;
            if (cudaVersion && cudaVersion !== 'null') {
                command = this.getClosestVersionCommand(cudaVersion, envName);
            } else {
                command = this.getClosestVersionCommand('cpu', envName);
            }
            console.log(command);
            if (this.remote_id === 'local') {
                await this.executeCommand(`${command}`);
            } else {
                await this.executeCommand(`${this.installDirs.conda}/bin/${command}`);
            }
            console.log(`PyTorch installed in environment ${envName} with command: ${command}`);
        } catch (error) {
            console.error('Error installing PyTorch:', error);
            return null;
        }
    }

    async installPaddle(envName, paddleVersion) {
        // console.log('installPaddle paddleVersion = ', paddleVersion);
        try {
            let cudaVersion = this.systemInfo.cudaInfo.nvccVersion;
            if (!cudaVersion || cudaVersion === 'null') {
                cudaVersion = this.systemInfo.cudaInfo.nvidiaSmiCUDA;
            }

            // 从配置中找到最合适的安装版本和命令
            const cudaMajorVersion = `CUDA${parseInt(cudaVersion.split('.')[0])}`;
            const paddleCandidates = this.paddleVersions[paddleVersion];
            // console.log('paddleCandidates = ', paddleCandidates);

            // 找到最接近的版本
            const compatibleVersions = Object.keys(paddleCandidates)
                .filter((v) => v.startsWith(cudaMajorVersion))
                .sort((a, b) => parseFloat(b.split('.')[1]) - parseFloat(a.split('.')[1]));

            let selectedVersion = null;
            for (let version of compatibleVersions) {
                if (parseFloat(version.split('.')[1]) <= parseFloat(cudaVersion.split('.')[1])) {
                    selectedVersion = version;
                    break;
                }
            }

            if (!selectedVersion) {
                console.error('No compatible CUDA version found.');
                return null;
            }

            const command = paddleCandidates[selectedVersion][0];
            console.log(command);

            // 开始执行安装命令
            if (this.remote_id === 'local') {
                await this.executeCommand(`${command}`);
            } else {
                await this.executeCommand(`${this.installDirs.conda}/bin/conda run -n ${envName} ${command}`);
            }
            console.log(`Paddle installed in environment ${envName} with command: ${command}`);
        } catch (error) {
            console.error('Error installing Paddle:', error);
            return null;
        }
    }

    async installPackages(envName, packages) {
        const packageList = Array.isArray(packages) ? packages.join(' ') : packages;
        const command = `${this.condaPath} install -y -n ${envName} ${packageList}`;
        try {
            await this.executeCommand(command);
            console.log(`Packages installed in environment ${envName}`);
        } catch (error) {
            console.error('Error installing packages:', error);
            return null;
        }
    }

    async installPackagesWithPip(envName, packages) {
        const packageList = Array.isArray(packages) ? packages.join(' ') : packages;
        const command = `${this.condaPath} run -n ${envName} pip install ${packageList} -i https://pypi.tuna.tsinghua.edu.cn/simple`;
        try {
            await this.executeCommand(command);
            console.log(`Packages installed via pip in environment ${envName}`);
        } catch (error) {
            console.error('Error installing packages via pip:', error);
            return null;
        }
    }

    async installFromRequirements(envName, requirementsPath) {
        const command = `${this.condaPath} install -y -n ${envName} --file ${requirementsPath}`;
        try {
            await this.executeCommand(command);
            console.log(`Packages from ${requirementsPath} installed in environment ${envName}`);
        } catch (error) {
            console.error('Error installing from requirements file:', error);
            return null;
        }
    }
}

module.exports = { CondaManager };
